# -*- coding:utf-8 -*-
"""
@author: fsksf

@since: 2022/4/9 13:55
"""
from typing import Union, List, Dict, Set, Optional, Awaitable, Any
import uuid
from datetime import datetime, date
import tornado.websocket
import tornado.escape
from vcat.server.handlers import *
import json
import enum
from typing import Union
from functools import wraps
import pyrestful.rest
import secrets
from cacheout import LRUCache

from .status_code import make_response, Code, WebSocketResp
from .logger import ser_logger

name_user_cache = LRUCache()


class CustomEncoder(json.JSONEncoder):
    def default(self, o):
        if o != o:
            return 'null'
        elif isinstance(o, datetime):
            return datetime.strftime(o, '%Y-%m-%d %H:%M:%S')
        elif isinstance(o, date):
            return datetime.strftime(o, '%Y-%m-%d')

        try:
            return super(CustomEncoder, self).default(o)
        except TypeError:
            if isinstance(o, enum.Enum):
                return o.value
            elif o != o:
                return 'null'
            return str(o)


class BaseResource(pyrestful.rest.RestHandler):
    """解决JS跨域请求问题"""
    def set_default_headers(self):
        self.set_header('Access-Control-Allow-Origin', '*')
        self.set_header('Access-Control-Allow-Headers', 'x-requested-with') # 后增加’ x-token‘
        self.set_header('Access-Control-Allow-Methods', 'POST, GET, OPTIONS, PUT, DELETE')
        self.set_header('Access-Control-Max-Age', 1000)
        self.set_header('Content-type', 'application/json; charset=UTF-8')
        self.set_header('Access-Control-Allow-Credentials', 'true')


    def options(self):
        self.set_status(200)
        self.finish()

    def get_current_user(self):
        _, user = self.get_curr_user()
        return user

    def get_curr_user(self):
        """
        获取当前user
        :return: (token是否有效， BaseUser对象)
        """
        user = None
        token_effective = False             # token是否有效
        token = self.get_token_arg()
        if token:
            username, _ = token.split('|')
            user: BaseUser = name_user_cache.get(username)
            if user and token == user.token:
                token_effective = True
        return token_effective, user

    @staticmethod
    def set_curr_user(user):
        """
        设置的当前用户
        :param user:
        :return: token
        """
        token = f'{user.username}|{secrets.token_urlsafe(30)}'
        user.token = token
        name_user_cache.set(user.username, user)
        return token

    def log_out_curr_user(self):
        user = self.get_current_user()
        name_user_cache.delete(user.username)

    def get_token_arg(self):
        return self.request.headers.get('X-Token')

    def write(self, chunk: Union[str, bytes, dict]) -> None:
        chunk = json.dumps(chunk, cls=CustomEncoder, allow_nan=False)
        super(BaseResource, self).write(chunk)


class Authenticated:
    def __init__(self, role_list: list):
        self.role_list = set(role_list)           # 允许角色列表

    def __call__(self, method):
        @wraps(method)
        def role_wrapper(_self: pyrestful.rest.RestHandler, *args, **kwargs):
            userinfo = _self.get_current_user()
            if userinfo:
                if self.role_list.intersection(userinfo.roles):
                    ser_logger.debug(f'user: {userinfo.name} authenticated pass')
                    return method(_self, *args, **kwargs)
                else:
                    ser_logger.info(f'{userinfo.name} permission_not_allow')
                    return make_response(Code.permission_not_allow)
            elif not AUTH:
                # 本地展示，可以开启不需要认证
                return method(_self, *args, **kwargs)
            else:
                ser_logger.debug('user: not login')
                return make_response(Code.user_not_login)
        return role_wrapper


class WebSocketHandler(tornado.websocket.WebSocketHandler):
    bao_cons: Set["WebSocketHandler"] = set()
    bao_waiters: Dict[str, "WebSocketHandler"] = {}
    con_key = None

    def set_default_headers(self):
        # self.set_header('Access-Control-Allow-Origin', '*')
        self.set_header('Access-Control-Allow-Methods', 'POST, GET, OPTIONS, PUT, DELETE')
        self.set_header('Access-Control-Max-Age', 1000)
        self.set_header('Content-type', 'application/json; charset=UTF-8')
        self.set_header('Access-Control-Allow-Credentials', 'true')
        self.set_header('Access-Control-Allow-Headers', 'Content-Type, Set-Cookie, x-token')    # 后增加’ x-token‘

    def open(self, *args: str, **kwargs: str) -> Optional[Awaitable[None]]:
        sole_id = str(uuid.uuid4()).upper()
        ser_logger.info(f'new ws {sole_id}')
        self.con_key = sole_id
        self.bao_waiters[sole_id] = self
        self.bao_cons.add(self)
        msg_resp = WebSocketResp(
            msg_type='open-websocket',
            msg_data=sole_id
        )
        self.write_message(msg_resp.to_dict())
        ser_logger.info('websocket opened!')

    def allow_draft76(self):
        return True

    def check_origin(self, origin: str):
        return True

    def on_close(self):
        self.bao_cons.remove(self)
        self.bao_waiters.pop(self.con_key)

        ser_logger.info("websocket closed!")
        print(self.bao_cons)

    def write_message(
        self, message: Union[bytes, str, Dict[str, Any]], binary: bool = False
    ) -> "Future[None]":
        message = json.dumps(message, cls=CustomEncoder, allow_nan=False)
        return super(WebSocketHandler, self).write_message(message=message, binary=binary)